#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

import os
import re
import json
import socket
import struct
import asyncio
import pathlib
import tempfile
import subprocess
from OpenSSL import SSL
from fm_util import Util
from fm_util import AvahiServiceBrowser
from fm_param import FmConst


class RemoteOrLocalBuild:

    class ExecContext:

        def __init__(self, parent, syncDownType=None):
            assert parent._execContext is None
            assert syncDownType in ["kernel", "system"]

            self._obj = parent
            self._syncDownType = syncDownType

            self._obj._execContext = self

        def dispose(self, bAborted=False):
            assert self._obj._execContext is self

            while self._obj.buildServer is not None:
                if self._syncDownType == "system":
                    if not bAborted:
                        self._obj._p.printInfo(">> Synchronizing down system files...")
                    else:
                        self._obj._p.printInfo(">> Aborted, synchronizing down system files...")
                    self._obj.buildServer.syncDownSystem()
                    print("")
                elif self._syncDownType == "kernel":
                    if not bAborted:
                        self._obj._p.printInfo(">> Synchronizing down /boot, /lib/modules and /lib/firmware...")
                    else:
                        self._obj._p.printInfo(">> Aborted, synchronizing down /boot, /lib/modules and /lib/firmware...")
                    self._obj.buildServer.syncDownKernel()
                    print("")
                else:
                    assert False
                break

            self._obj._execContext = None

        def exec(self, *args):
            if self._obj.buildServer is None:
                Util.cmdExec(*args)
            else:
                self._obj.buildServer.sshExec(*args)

        def getFile(self, resultFile, encoding):
            if self._obj.buildServer is None:
                with open(resultFile, "r", encoding=encoding) as f:
                    return f.read()
            else:
                return self._obj.buildServer.getFile(resultFile).decode(encoding)

        def __enter__(self):
            return self

        def __exit__(self, type, value, traceback):
            # bAborted is True when there's exception
            self.dispose(bAborted=(type is not None))

    def __init__(self, infoPrinter):
        self._p = infoPrinter
        self._buildServer = None
        self._execContext = None

        if not os.path.exists(FmConst.buildServerConfFile):
            return

        self._p.printInfo(">> Selecting build server...")
        self._buildServer = self._getBuildServer()
        print("")

        if self._buildServer is None:
            return

        # sync up and start working
        self._p.printInfo(">> Synchronizing up...")
        self._buildServer.syncUp()
        self._buildServer.startWorking()
        print("")

    @property
    def buildServer(self):
        return self._buildServer

    def getExecContext(self, syncDownType=None):
        return self.ExecContext(self, syncDownType)

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        assert self._execContext is None
        del self._execContext

        if self._buildServer is not None:
            self._buildServer.dispose()
        del self._buildServer

    @classmethod
    def _getBuildServer(cls):
        bForce, bAutoDiscover, serverList = cls._getVariables()

        ret = cls._getBuildServerImpl(bForce, bAutoDiscover, serverList)
        if ret is None:
            ss = "No build server defined, fallback to normal process."
            if bAutoDiscover:
                ss = ss.replace("d,", "d or discovered,")
            print(ss)
            return None
        name, server, bForce = ret

        if ":" in server:
            ip = server.split(":")[0]
            ctrlPort = int(server.split(":")[1])
        else:
            ip = server
            ctrlPort = FmConst.buildServerDefaultPort

        ret = RemoteBuildServer(ip, ctrlPort)
        ret.checkCertKey()
        try:
            ret.connectAndInit()
            print("Use build server \"%s\"." % (name))
            return ret
        except:
            if not bForce:
                print("Failed to use build server \"%s\", fallback to normal process." % (name))
                return None                             # fail
            else:
                raise                                   # fail

    @staticmethod
    def _getVariables():
        # get variables
        bForce = False
        bAutoDiscover = False
        serverList = []
        with open(FmConst.buildServerConfFile, 'r') as f:
            lineList = f.read().split()
            for i in range(0, len(lineList)):
                line = lineList[i].strip()
                if line == "" or line.startswith("#"):
                    continue

                m = re.search("^FORCE=(.*)$", line)
                if m is not None and m.group(1).lower() in ["true", "1"]:
                    bForce = True
                    continue

                m = re.search("^AUTO_DISCOVER=(.*)$", line)
                if m is not None and m.group(1).lower() in ["true", "1"]:
                    bAutoDiscover = True
                    continue

                m = re.search("^BUILD_SERVER=(.*)$", line)
                if m is not None:
                    serverList = m.group(1).split()
                    continue

                raise Exception("invalid content at line %d in %s" % (i + 1, FmConst.buildServerConfFile))

        return (bForce, bAutoDiscover, serverList)

    @staticmethod
    def _getBuildServerImpl(bForce, bAutoDiscover, serverList):
        # we prefer static configured build servers
        if serverList != []:
            return (serverList[0], serverList[0], bForce)

        # discover build server
        if bAutoDiscover:
            browser = AvahiServiceBrowser("_syncup_gentoo._tcp")
            browser.run()
            for name, addr, port in browser.get_result_list():
                return (name, "%s:%d" % (addr, port), bForce)

        # no build server found
        if bForce:
            raise Exception("no build server available")
        return None


class RemoteBuildServer:

    def __init__(self, server, ctrlPort):
        self._hostname = server
        self._ctrlPort = ctrlPort

        self._sock = None
        self._sslSock = None

        self._wSshPort = None
        self._wRsyncPort = None
        self._wCatFilePort = None
        self._identityFile = None
        self._cfgFile = None

        self._asyncJobCount = 0

    def checkCertKey(self):
        if not os.path.exists(FmConst.myCertFile) or not os.path.exists(FmConst.myPrivKeyFile):
            raise Exception("%s or %s does not exist" % (FmConst.myCertFile, FmConst.myPrivKeyFile))

    def connectAndInit(self):
        try:
            self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self._sock.connect((self._hostname, self._ctrlPort))

            ctx = SSL.Context(SSL.TLSv1_2_METHOD)
            ctx.use_certificate_file(FmConst.myCertFile)
            ctx.use_privatekey_file(FmConst.myPrivKeyFile)
            self._sslSock = SSL.Connection(ctx, self._sock)
            self._sslSock.set_connect_state()

            self._sendRequestObj({
                "command": "init",
                "hostname": socket.gethostname(),
                "cpu-arch": Util.getCpuArch(),
                "cpu-model": Util.getCpuModel(),
                "plugin": "gentoo",
            })
            resp = self._recvReponseObj()
            if "error" in resp:
                raise Exception(resp["error"])
        except:
            self.dispose()
            raise

    def dispose(self):
        assert self._asyncJobCount == 0

        if self._identityFile is not None:
            os.unlink(self._identityFile)
            self._identityFile = None
        if self._cfgFile is not None:
            os.unlink(self._cfgFile)
            self._cfgFile = None
        self._wSshPort = None
        self._wRsyncPort = None
        self._wCatFilePort = None

        if self._sslSock is not None:
            try:
                self._sslSock.shutdown()
            except SSL.Error:
                # ignore all kinds of error
                pass
            self._sslSock = None
        if self._sock is not None:
            self._sock.close()
            self._sock = None

        self._ctrlPort = None
        self._hostname = None

    def syncUp(self):
        # enter stage
        self._sendRequestObj({
            "command": "stage-syncup",
        })
        resp = self._recvReponseObj()
        if "error" in resp:
            raise Exception(resp["error"])
        assert resp["return"]["stage"] == "syncup"

        # rsync
        stunnelCfgFile, newPort, proc = self._createStunnelProcess(resp["return"]["rsync-port"])
        try:
            cmd = ""
            cmd += "rsync -a -z -hhh --delete --delete-excluded --partial --info=progress2 "
            for fn in self._ignoredPatternsWhenSyncUp():
                cmd += "-f '- %s' " % (fn)
            cmd += "-f '+ /bin' "                                       # /bin may be a symlink or directory
            cmd += "-f '+ /bin/***' "
            cmd += "-f '+ /boot/***' "
            cmd += "-f '+ /etc/***' "
            cmd += "-f '+ /lib' "                                       # /lib may be a symlink or directory
            cmd += "-f '+ /lib/***' "
            cmd += "-f '+ /lib32' "                                     # /lib32 may be a symlink or directory
            cmd += "-f '+ /lib32/***' "
            cmd += "-f '+ /lib64' "                                     # /lib64 may be a symlink or directory
            cmd += "-f '+ /lib64/***' "
            cmd += "-f '+ /opt/***' "
            cmd += "-f '+ /sbin' "                                      # /sbin may be a symlink or directory
            cmd += "-f '+ /sbin/***' "
            cmd += "-f '+ /usr/***' "
            cmd += "-f '+ /var' "
            cmd += "-f '+ /var/cache' "
            cmd += "-f '+ /var/cache/edb/***' "
            cmd += "-f '+ /var/cache/portage/***' "
            cmd += "-f '+ /var/db' "
            cmd += "-f '+ /var/db/pkg/***' "
            cmd += "-f '+ /var/lib' "
            cmd += "-f '+ /var/lib/portage/***' "
            cmd += "-f '- /**' "
            cmd += "/ rsync://127.0.0.1:%d/main" % (newPort)
            Util.shellExec(cmd)
        finally:
            proc.terminate()
            proc.wait()
            os.unlink(stunnelCfgFile)

    def startWorking(self):
        self._sendRequestObj({
            "command": "stage-working",
        })
        resp = self._recvReponseObj()
        if "error" in resp:
            raise Exception(resp["error"])
        assert resp["return"]["stage"] == "working"

        self._wSshPort = resp["return"]["ssh-port"]
        self._wRsyncPort = resp["return"]["rsync-port"]
        self._wCatFilePort = resp["return"]["catfile-port"]

        self._identityFile = None
        if True:
            fullfn = tempfile.mktemp()
            pathlib.Path(fullfn).write_text("")
            os.chmod(fullfn, 0o600)
            pathlib.Path(fullfn).write_text(resp["return"]["ssh-key"])
            self._identityFile = fullfn

        self._cfgFile = tempfile.mktemp()
        if True:
            buf = ""
            buf += "LogLevel QUIET\n"
            buf += "\n"
            buf += "KbdInteractiveAuthentication no\n"
            buf += "PasswordAuthentication no\n"
            buf += "PubkeyAuthentication yes\n"
            buf += "PreferredAuthentications publickey\n"
            buf += "\n"
            buf += "IdentityFile %s\n" % (self._identityFile)
            buf += "UserKnownHostsFile /dev/null\n"
            buf += "StrictHostKeyChecking no\n"
            buf += "\n"
            buf += "SendEnv LANG LC_*\n"
            with open(self._cfgFile, "w") as f:
                f.write(buf)

    def sshExec(self, *args, base64=False):
        assert self._wSshPort is not None

        # FIXME: ssh sucks that it must use a shell to execute remote command
        if base64:
            args2 = list(args)
            for i in range(1, len(args2)):
                args2[i] = base64.b64encode(args2[i].encode("ascii")).decode("ascii")
        else:
            args2 = args

        # "-t" can get Ctrl+C controls remote process
        # XXXXX so that we forward signal to remote process, FIXME
        cmd = "ssh -t -e none -p %d -F %s %s %s" % (self._wSshPort, self._cfgFile, self._hostname, " ".join(args2))
        Util.shellExec(cmd)

    async def asyncStartSshExec(self, cmd, *kargs, loop=None):
        assert self._wSshPort is not None
        assert loop is not None

        cmd = "ssh -t -e none -p %d -F %s %s %s %s" % (self._wSshPort, self._cfgFile, self._hostname, cmd, " ".join(kargs))
        proc = await asyncio.create_subprocess_shell(cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT, loop=loop)
        self._asyncJobCount += 1
        return (proc, proc.stdout)

    async def asyncWaitSshExec(self, proc):
        try:
            retcode = await proc.wait()
            if retcode != 0:
                raise subprocess.CalledProcessError(retcode, proc.args)      # use subprocess.CalledProcessError since there's no equivalent in asyncio
        finally:
            self._asyncJobCount -= 1

    def getFile(self, filename):
        assert self._wSshPort is not None

        stunnelCfgFile, newPort, proc = self._createStunnelProcess(self._wCatFilePort)
        try:
            _sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            try:
                _sock.connect(("127.0.0.1", newPort))
                _sock.sendall(struct.pack("!I", len(filename.encode("utf-8"))))
                _sock.sendall(filename.encode("utf-8"))
                errCode = struct.unpack("!c", self._sockRecvAll(_sock, 1))[0]
                dataLen = struct.unpack("!Q", self._sockRecvAll(_sock, 8))[0]
                data = self._sockRecvAll(_sock, dataLen)
                if errCode != b'\x00':
                    raise Exception(data.decode("utf-8"))
                return data
            finally:
                _sock.close()
        finally:
            proc.terminate()
            proc.wait()
            os.unlink(stunnelCfgFile)

    def syncDownKernel(self):
        assert self._wSshPort is not None

        stunnelCfgFile, newPort, proc = self._createStunnelProcess(self._wRsyncPort)
        try:
            cmd = ""
            cmd += "rsync -a -z -hhh --delete --info=progress2 "
            cmd += "-f '+ /boot' "
            cmd += "-f '+ /boot/config-*' "
            cmd += "-f '+ /boot/initramfs-*' "
            cmd += "-f '+ /boot/kernel-*' "
            cmd += "-f '+ /boot/System.map-*' "
            cmd += "-f '+ /boot/history/***' "
            cmd += "-f '+ /lib' "
            cmd += "-f '+ /lib/modules/***' "
            cmd += "-f '+ /lib/firmware/***' "
            cmd += "-f '- /**' "
            cmd += "rsync://127.0.0.1:%d/main /" % (newPort)
            Util.shellExec(cmd)
        finally:
            proc.terminate()
            proc.wait()
            os.unlink(stunnelCfgFile)

    def syncDownSystem(self):
        assert self._wSshPort is not None

        stunnelCfgFile, newPort, proc = self._createStunnelProcess(self._wRsyncPort)
        try:
            cmd = ""
            cmd += "rsync -a -z -hhh --delete --info=progress2 "
            for fn in self._ignoredPatternsWhenSyncDown():
                cmd += "-f '- %s' " % (fn)
            cmd += "-f '+ /bin' "                                       # /bin may be a symlink or directory
            cmd += "-f '+ /bin/***' "
            cmd += "-f '+ /etc/***' "
            cmd += "-f '+ /lib' "                                       # /lib may be a symlink or directory
            cmd += "-f '+ /lib/***' "
            cmd += "-f '+ /lib32' "                                     # /lib may be a symlink or directory
            cmd += "-f '+ /lib32/***' "
            cmd += "-f '+ /lib64' "                                     # /lib may be a symlink or directory
            cmd += "-f '+ /lib64/***' "
            cmd += "-f '+ /opt/***' "
            cmd += "-f '+ /sbin' "                                      # /sbin may be a symlink or directory
            cmd += "-f '+ /sbin/***' "
            cmd += "-f '+ /usr/***' "
            cmd += "-f '+ /var' "
            cmd += "-f '+ /var/cache' "
            cmd += "-f '+ /var/cache/edb/***' "
            cmd += "-f '+ /var/cache/portage/***' "
            cmd += "-f '+ /var/db' "
            cmd += "-f '+ /var/db/pkg/***' "
            cmd += "-f '+ /var/lib' "
            cmd += "-f '+ /var/lib/portage/***' "
            cmd += "-f '- /**' "
            cmd += "rsync://127.0.0.1:%d/main /" % (newPort)
            Util.shellExec(cmd)
        finally:
            proc.terminate()
            proc.wait()
            os.unlink(stunnelCfgFile)

    def syncDownDirectory(self, dirname, quiet=False):
        assert self._wSshPort is not None
        assert dirname.startswith("/")

        dirname = os.path.realpath(dirname)
        stunnelCfgFile, newPort, proc = self._createStunnelProcess(self._wRsyncPort)
        try:
            cmd = ""
            cmd += "rsync -a -z -hhh --delete %s " % ("--quiet" if quiet else "--info=progress2")
            for fn in self._ignoredPatternsWhenSyncDown():
                cmd += "-f '- %s' " % (fn)
            if True:
                buf = "-f '+ %s/***' " % (dirname)
                dirname = os.path.dirname(dirname)
                while dirname != "/":
                    buf = "-f '+ %s' " % (dirname) + buf
                    dirname = os.path.dirname(dirname)
                cmd += buf
            cmd += "-f '- /**' "
            cmd += "rsync://127.0.0.1:%d/main /" % (newPort)
            Util.shellExec(cmd)
        finally:
            proc.terminate()
            proc.wait()
            os.unlink(stunnelCfgFile)

    def syncDownWildcardList(self, wildcardList, quiet=False):
        assert self._wSshPort is not None
        assert [x.startswith("/") for x in wildcardList]

        stunnelCfgFile, newPort, proc = self._createStunnelProcess(self._wRsyncPort)
        try:
            cmd = ""
            cmd += "rsync -a -z -hhh --delete %s " % ("--quiet" if quiet else "--info=progress2")
            for fn in self._ignoredPatternsWhenSyncDown():
                cmd += "-f '- %s' " % (fn)
            for wildcard in wildcardList:
                buf = "-f '+ %s' " % (wildcard)
                dirname = os.path.dirname(wildcard)
                while dirname != "/":
                    buf = "-f '+ %s' " % (dirname) + buf
                    dirname = os.path.dirname(dirname)
                cmd += buf
            cmd += "-f '- /**' "
            cmd += "rsync://127.0.0.1:%d/main /" % (newPort)
            Util.shellExec(cmd)
        finally:
            proc.terminate()
            proc.wait()
            os.unlink(stunnelCfgFile)

    def _ignoredPatternsWhenSyncUp(self):
        return [
            FmConst.buildServerConfFile,
            FmConst.myCertFile,
            FmConst.myPrivKeyFile,
        ]

    def _ignoredPatternsWhenSyncDown(self):
        return self._ignoredPatternsWhenSyncUp() + [
            FmConst.portageCfgMakeConf,
            FmConst.portageMirrorsFile,
            "/etc/resolv.conf",
        ]

    def _createStunnelProcess(self, port):
        stunnelCfgFile = tempfile.mktemp()
        newPort = Util.getFreeTcpPort()
        try:
            buf = ""
            buf += "cert = %s\n" % (FmConst.myCertFile)
            buf += "key = %s\n" % (FmConst.myPrivKeyFile)
            buf += "\n"
            buf += "client = yes\n"
            buf += "foreground = yes\n"
            buf += "\n"
            buf += "[rsync]\n"
            buf += "accept = %d\n" % (newPort)
            buf += "connect = %s:%d\n" % (self._hostname, port)
            with open(stunnelCfgFile, "w") as f:
                f.write(buf)

            proc = subprocess.Popen("stunnel %s 2>/dev/null" % (stunnelCfgFile), shell=True)
            Util.waitTcpService("0.0.0.0", newPort)

            return (stunnelCfgFile, newPort, proc)
        except:
            os.unlink(stunnelCfgFile)
            raise

    def _sendRequestObj(self, requestObj):
        s = json.dumps(requestObj) + "\n"
        self._sslSock.send(s.encode("iso8859-1"))

    def _recvReponseObj(self):
        buf = b''
        while True:
            buf += self._sslSock.recv(4096)
            i = buf.find(b'\n')
            if i >= 0:
                assert i == len(buf) - 1
                return json.loads(buf[:i].decode("iso8859-1"))

    def _sockRecvAll(self, _sock, datasize):
        buf = b''
        while len(buf) < datasize:
            buf2 = _sock.recv(datasize - len(buf))
            if len(buf2) == 0:
                raise EOFError()
            buf += buf2
        return buf
